from fastapi import Header, HTTPException
from app.core.config import settings
import hmac
import hashlib
from datetime import datetime, timedelta
import json

def generate_token(source: str) -> str:
    """生成token
    
    Args:
        source: 数据来源标识
        
    Returns:
        str: 生成的token
    """
    timestamp = datetime.now().timestamp()
    message = f"{source}:{timestamp}"
    signature = hmac.new(
        settings.TOKEN_SECRET.encode(),
        message.encode(),
        hashlib.sha256
    ).hexdigest()
    return f"{message}:{signature}"

def verify_token_signature(token: str) -> bool:
    """验证token签名
    
    Args:
        token: 待验证的token
        
    Returns:
        bool: 验证是否通过
    """
    try:
        if token in settings.ALLOWED_TOKENS:  # 支持预设token
            return True
            
        message, signature = token.rsplit(":", 1)
        expected_signature = hmac.new(
            settings.TOKEN_SECRET.encode(),
            message.encode(),
            hashlib.sha256
        ).hexdigest()
        return hmac.compare_digest(signature, expected_signature)
    except Exception as e:
        print(f"Token验证错误: {str(e)}")
        return False

async def verify_token(token: str = Header(..., description="认证令牌")) -> str:
    """Token验证中间件
    
    Args:
        token: 请求头中的token
        
    Returns:
        str: 验证通过的token
        
    Raises:
        HTTPException: token验证失败时抛出
    """
    if not token:
        raise HTTPException(status_code=401, detail="Token未提供")
    
    if not verify_token_signature(token):
        raise HTTPException(status_code=401, detail="无效的Token")
    
    try:
        if token in settings.ALLOWED_TOKENS:  # 预设token直接通过
            return token
            
        # 动态token需要验证时效性
        message = token.split(":")[1]
        timestamp = float(message)
        token_time = datetime.fromtimestamp(timestamp)
        if datetime.now() - token_time > timedelta(minutes=settings.TOKEN_EXPIRE_MINUTES):
            raise HTTPException(status_code=401, detail="Token已过期")
    except ValueError:
        raise HTTPException(status_code=401, detail="无效的Token格式")
    
    return token

async def verify_source(source: str = Header(..., description="数据来源")) -> str:
    """验证数据来源
    
    Args:
        source: 请求头中的数据来源标识
        
    Returns:
        str: 验证通过的数据来源
        
    Raises:
        HTTPException: 数据来源验证失败时抛出
    """
    if source not in settings.ALLOWED_SOURCES:
        raise HTTPException(status_code=400, detail="无效的数据来源")
    return source